use crate::vector::{dot_product, Vector};
use anyhow::anyhow;
use std::fmt::{Debug, Display, Formatter};
use std::ops::{Add, AddAssign, Mul, MulAssign};
use std::sync::mpsc;
use std::thread;

const NUM: i32 = 4;

pub struct Matrix<T: Debug> {
    data: Vec<T>,
    row: usize,
    col: usize,
}

pub struct MsgInput<T> {
    idx: usize,
    row: Vector<T>,
    col: Vector<T>,
}

pub struct MsgOutput<T> {
    idx: usize,
    value: T,
}

pub struct Msg<T> {
    input: MsgInput<T>,
    sender: oneshot::Sender<MsgOutput<T>>,
}

impl<T> Mul for Matrix<T>
where
    T: Add<Output = T>
        + Default
        + AddAssign
        + Copy
        + MulAssign
        + Mul<Output = T>
        + Debug
        + Send
        + Sync
        + 'static,
{
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        multiply(&self, &rhs).expect("Failed to multiply")
    }
}

//使用T:Default
pub fn multiply<T>(a: &Matrix<T>, b: &Matrix<T>) -> anyhow::Result<Matrix<T>>
where
    T: Add<Output = T>
        + Default
        + AddAssign
        + Copy
        + MulAssign
        + Mul<Output = T>
        + Debug
        + Send
        + Sync
        + 'static,
{
    if a.col != b.row {
        return Err(anyhow!("Matrix dimensions do not match"));
    }
    let mut data = vec![T::default(); a.row * b.col];

    let senders = (0..NUM)
        .map(|_| {
            let (tx, rx) = mpsc::channel::<Msg<T>>();
            thread::spawn(move || {
                for msg in rx {
                    let value = dot_product(msg.input.row, msg.input.col)?;
                    if let Err(e) = msg.sender.send(MsgOutput {
                        idx: msg.input.idx,
                        value,
                    }) {
                        eprintln!("Error sending msg output: {e}");
                    }
                }
                Ok::<_, anyhow::Error>(())
            });
            tx
        })
        .collect::<Vec<_>>();

    //generate 4 thread which receive msg and do dot product

    let mut receivers = Vec::with_capacity(a.row * b.col);
    for i in 0..a.row {
        for j in 0..b.col {
            let row = Vector::new(&a.data[i * a.col..(i + 1) * a.col]);
            let col_data = b.data[j..]
                .iter()
                .step_by(b.col)
                .copied()
                .collect::<Vec<_>>();
            let col = Vector::new(col_data);
            let idx = i * b.col + j;
            let input = MsgInput { idx, row, col };
            let (tx, rx) = oneshot::channel();
            let msg = Msg { input, sender: tx };
            if let Err(e) = senders[idx % NUM as usize].send(msg) {
                eprintln!("Send error:{e:?}");
            }
            // data[i * b.col + j] += dot_product(row, col)?;
            // let ret = rx.recv()?;
            // data[ret.idx]  =ret.value;
            receivers.push(rx);
        }
    }

    for x in receivers {
        let ret = x.recv()?;
        data[ret.idx] = ret.value;
    }

    Ok(Matrix {
        data,
        row: a.row,
        col: b.col,
    })
}

impl<T: Debug> Matrix<T> {
    //impl Into<Vec<T>> 任何数据结构只有能convert T就能使用
    pub fn new(data: impl Into<Vec<T>>, row: usize, col: usize) -> Matrix<T> {
        Self {
            data: data.into(),
            row,
            col,
        }
    }
}

impl<T> Display for Matrix<T>
where
    T: Debug,
{
    //display a 2x3 as {1 2 3 , 4 5 6} 3x2 as {1 2, 3 4,5 6}
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for i in 0..self.row {
            for j in 0..self.col {
                write!(f, "{:?}", self.data[i * self.col + j])?;

                if j != self.col - 1 {
                    write!(f, " ")?;
                }
            }
            if i != self.row - 1 {
                write!(f, ",")?;
            }
        }
        Ok(())
    }
}

impl<T> Debug for Matrix<T>
where
    T: Display + Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Matrix(row={},col={},{{ {} }})",
            self.row, self.col, self
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /***
    1 2 3     1*1 + 2*3 + 3*5 15+6
    4 5 6


    1 2
    3 4
    5 6
    */

    #[test]
    fn test_matrix_multiply() {
        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
        let b = Matrix::new([1, 2, 3, 4, 5, 6], 3, 2);
        let c = a * b;
        println!("{c}");
        assert_eq!(format!("{c:?}"), "Matrix(row=2,col=2,{ 22 28,49 64 })");
    }
}
